from casetoolbox.casestab import casestab
from casetoolbox.casestab import math_functions as mf
from casetoolbox.casestab import model_assembler as ma
import os
import sys
import json
from copy import deepcopy
import numpy as np
from matplotlib import pyplot as plt

# Set Python path to the file's directory
script_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, script_dir)
os.chdir(script_dir)

# Get current work directory
work_folder = os.getcwd()
# Set result folder
res_folder = 'SteadyState/results_6x6'

# =============================================================================
# Create rotor models with monopole substructure
# =============================================================================
# Read json file        
with open(os.path.join(work_folder,'DTU10MW_6x6.json')) as json_file:
    input_para = json.load(json_file)

# Define substructures =======================================================
subs_para_set=[]

# MONOPOLE Type and base data of substructure (connected to ground at bottom)
monopole_para={}
monopole_para['znode']=np.linspace(-115.63, 0, 11)  # Tower height from HAWC2 data: bottom to top
monopole_para['type']='HAWC2elements'
monopole_para['name']=os.path.join(work_folder,'HAWC2_beam_dtu10mw.dat')
monopole_para['bname']='tower'
monopole_para['norder']=0
subs_para={}
subs_para['isubs']=0
subs_para['name']='monopole'
subs_para['Sbase']=np.eye(3)
subs_para['bearing']=''
subs_para['type']='corotbeam'
subs_para['isubs_connection'] = -1  # Connected to ground
subs_para['inode_connection'] = -1  # Node 0 (bottom) is fixed to ground
subs_para['para']=monopole_para
subs_para_set.append(subs_para)

# HUB Type and base data of substructure (connected to tower top)
rigid_para={}
rigid_para['nodes']=np.array([[0.0,0.0],[0.0,0.0],[0.0,input_para['hub']['radius']]])
subs_para={}
subs_para['isubs']=1
subs_para['name']='hub1'
subs_para['Sbase']=np.eye(3)
subs_para['bearing']='constant_speed 2 rotor_speed rpm'
subs_para['type']='rigid_vector'
subs_para['isubs_connection'] = 0  # Connected to monopole/tower (substructure 0)
subs_para['inode_connection'] = 10  # Tower top node (z=0, node index 10 of 11 nodes, where node 0 is at bottom)
subs_para['para']=rigid_para
subs_para_set.append(subs_para)

# ROTOR Type and base data of substructure (connected to hub)
corotbeam_para={}
if input_para['blade']['structure']['type'] == 'HAWC2elements':
    corotbeam_para['znode']=input_para['blade']['structure']['znode']
    corotbeam_para['type']=input_para['blade']['structure']['type']
    corotbeam_para['name']=input_para['blade']['structure']['name']
    corotbeam_para['bname'] = input_para['blade']['structure']['bname']
    corotbeam_para['norder']=input_para['blade']['structure']['norder']
else:
    corotbeam_para['znode']=input_para['blade']['structure']['znode']
    corotbeam_para['type']=''
    corotbeam_para['name']=input_para['blade']['structure']['st_file']
    corotbeam_para['setno']=input_para['blade']['structure']['setno']
    corotbeam_para['subsetno']=input_para['blade']['structure']['subsetno']
    corotbeam_para['nintp']=input_para['blade']['structure']['nintp']
    corotbeam_para['norder']=input_para['blade']['structure']['norder']
subs_para={}
subs_para['isubs']=2
subs_para['name']='blade1'
subs_para['Sbase']=mf.R1(np.radians(input_para['hub']['cone'])) # Coning
subs_para['bearing']='constant_angle -3 pitch deg'
subs_para['type']='corotbeam'
subs_para['isubs_connection'] = 1  # Connected to hub (substructure 1)
subs_para['inode_connection'] = 1  # Last node of hub
subs_para['para']=corotbeam_para
subs_para_set.append(subs_para)

# Define blades ==============================================================
blade_para_set=[]
blade_para={}
blade_para['geo_file']=input_para['blade']['aerodynamics']['ae_file']
blade_para['pro_file']=input_para['blade']['aerodynamics']['pc_file']
blade_para['geo_set']=input_para['blade']['aerodynamics']['ae_setno']
blade_para['zaero'] = input_para['blade']['aerodynamics']['zaero']
blade_para['geo_inter'] = input_para['blade']['aerodynamics']['geo_inter']
blade_para['ae_inter'] = input_para['blade']['aerodynamics']['ae_inter']
blade_para['pro_inter'] = input_para['blade']['aerodynamics']['pc_inter']
blade_para['substructures']=[2]  # Blade uses substructure 2
blade_para_set.append(blade_para)

# Define wake
wake_para_set=[]
wake_para={}
wake_para['type']=input_para['wake']['type']
wake_para['a_of_CT_model']=input_para['wake']['a_of_CT_model']
wake_para['tip_correction']=input_para['wake']['tip_correction']
wake_para['number_of_blades']=input_para['rotor']['number_of_blades']
wake_para_set.append(wake_para)

# Define wind
wind_para_set=[]
wind_para={}
wind_para['windtype']='uniform'
wind_para['umean']=0.0
wind_para['density'] = input_para['wind']['density']
wind_para_set.append(wind_para)

# Define axissymmetric rotor
rotor_para_set=[]
rotor_para={}
rotor_para['isubs_rotorcenter']=1  # Rotor center at hub (substructure 1)
rotor_para['iaxis_rotorcenter']=2
rotor_para['type'] = 'axissym'
rotor_para['blades'] = [0]
rotor_para['blades_isubs'] = [2]
rotor_para['number_of_blades']=input_para['rotor']['number_of_blades']
rotor_para['iwake']=0
rotor_para['iwind']=0
rotor_para_set.append(rotor_para)

# Deflection flag
if 'deflection' in input_para.keys():
    include_deflection = (input_para['deflection'] == 1)
else:
    include_deflection = True

# Assemble model
model = ma.model(subs_para_set,blade_para_set,rotor_para_set,wake_para_set,wind_para_set)

# Create rotor_models-like object for compatibility
class rotor_models_with_monopole:
    def __init__(self, model, input_para, work_folder):
        self.folder = work_folder
        self.models = []
        if input_para['operation'].split()[0] == 'ops_file':
            self.ops = np.loadtxt(input_para['operation'].split()[1],skiprows=1)
            self.Nops = np.size(self.ops,axis=0)
            self.pwr = np.zeros(self.Nops)
            self.thr = np.zeros(self.Nops)
            self.CP  = np.zeros(self.Nops)
            for iops in range(self.Nops):
                model_copy = deepcopy(model)
                model_copy.substructures[1].bearing.state.speed = self.ops[iops,2]*np.pi/30.0
                model_copy.substructures[2].bearing.state.angle = np.radians(self.ops[iops,1])
                model_copy.rotors[0].wind.lookup.umean = self.ops[iops,0]
                self.models.append(model_copy)
    
    def steady_state_computation(self,ops_list=[]):
        if len(ops_list) == 0:
            ops_list = np.arange(self.Nops)
        for iops in ops_list:
            print('======================================================================================================')
            print('=========== Computing steady state for operation point {:2d} with wind speed {:4.1f} m/s ==================='.format(iops+1,self.ops[iops,0]))
            print('======================================================================================================')
            self.models[iops].compute_rotor_stationary_steady_state(0,1.0,include_deflection)
            self.pwr[iops] = self.models[iops].rotors[0].power
            self.thr[iops] = self.models[iops].rotors[0].thrust
            self.CP[iops]  = self.models[iops].rotors[0].CP
    
    def save_steady_state_results(self,prefix=''):
        defl_header = ['z-coord. [m] 1','ux_blade [m] 2','uy_blade [m] 3','uz_blade [m] 4','ux_rotor [m] 5','uy_rotor [m] 6','uz_rotor [m] 7','rotx_blade [deg] 8','roty_blade [deg] 9','rotz_blade [deg] 10']
        bem_header  = ['z-coord. [m] 1','a [-] 2','ap [-] 3','AoA [deg] 4', 'Urel [m/s] 5','Inplane Fx [N/m] 6','Axial Fy [N/m] 7','Moment [Nm/m] 8', 
                       'CL [-] 9','CD [-] 10','CM [-] 11','CLp [1/rad] 12','CDp [1/rad] 13','vx [m/s] 14','vy [m/s] 15','CT [-] 16','CQ [-] 17']
        for iops in range(self.Nops):
            # Structural results
            defl_filename = self.folder + prefix + 'defl_wsp{:d}mms_pitch{:d}mdeg_spd{:d}mrpm'.format(int(1e3*self.ops[iops,0]),int(1e3*self.ops[iops,1]),int(1e3*self.ops[iops,2]))
            # Collect data
            data = self.models[iops].substructures[2].create_data_for_deflection_state()
            # Save data
            header_txt='# Steady state results for all structural nodes \n'+''.join('{:>20s} '.format(text) for text in defl_header)
            np.savetxt(defl_filename,data,fmt='%20.10e',header=header_txt,comments='',delimiter=' ')
            # Aerdynamic BEM results
            bem_filename  = self.folder + prefix +  'BEM_wsp{:d}mms_pitch{:d}mdeg_spd{:d}mrpm'.format(int(1e3*self.ops[iops,0]),int(1e3*self.ops[iops,1]),int(1e3*self.ops[iops,2]))
            # Create data for output
            data = self.models[iops].rotors[0].create_data_for_BEM_results()
            # Save results 
            header_txt='# Steady state results for all aerodynamic calculation points \n'+''.join('{:>20s} '.format(text) for text in bem_header)
            np.savetxt(bem_filename,data,fmt='%20.10e',header=header_txt,comments='',delimiter=' ')

# Setup rotor models for each operational point (with monopole)
rotor_models = rotor_models_with_monopole(model, input_para, work_folder)
# Compute the steady for all operational points
rotor_models.steady_state_computation()
# Save the steady for all operational points
rotor_models.save_steady_state_results('/' + res_folder + '/ind/')
# Copy for shorter notation in plotting below
ops = rotor_models.ops
pwr = rotor_models.pwr
thr = rotor_models.thr
models = rotor_models.models
# Plot comparison to HS2 for each operational point
"""
for iops in range(len(models)):
    # Generate output data
    # Note: substructure[0]=monopole, [1]=hub, [2]=blade
    defl_data = models[iops].substructures[2].create_data_for_deflection_state()
    bem_data = models[iops].rotors[0].create_data_for_BEM_results()
    # Wind speed for labels
    wsp = ops[iops,0]
    # Read HS2 results
    hs2_ind = np.loadtxt(os.path.join(work_folder,'HAWCStab2_results/dtu10mw_u{:d}.ind'.format(int(1000*wsp))),skiprows = 1)
    # Plot deflections and forces in the rotor plane
    fig1=plt.figure(figsize=[12,6])
    axs1=fig1.subplots(3,2)
    axs1_2 = []
    for i in range(3):
        axdummy = []
        for j in range(2):
            axdummy.append(axs1[i,j].twinx())
        axs1_2.append(axdummy)
    axs1[0,0].plot(defl_data[:,0],defl_data[:,4],'.-',hs2_ind[:,0],hs2_ind[:,9],'.-')
    axs1[1,0].plot(defl_data[:,0],defl_data[:,5],'.-',hs2_ind[:,0],hs2_ind[:,10],'.-')
    axs1[2,0].plot(defl_data[:,0],defl_data[:,6],'.-',hs2_ind[:,0],hs2_ind[:,11],'.-')
    axs1[0,1].plot(defl_data[:,0],defl_data[:,9],'.-',hs2_ind[:,0],np.degrees(hs2_ind[:,34]*hs2_ind[:,37]),'.-')
    axs1[1,1].plot(bem_data[:,0],bem_data[:,5]*1e-3,'.-',hs2_ind[:,0],1e-3*hs2_ind[:,6],'.-')
    axs1[2,1].plot(bem_data[:,0],bem_data[:,6]*1e-3,'.-',hs2_ind[:,0],1e-3*hs2_ind[:,7],'.-')
    axs1[0,0].set_ylabel('x deflection [m]')
    axs1[1,0].set_ylabel('y deflection [m]')
    axs1[2,0].set_ylabel('z deflection [m]')
    axs1[0,1].set_ylabel('z rotation [deg]')
    axs1[1,1].set_ylabel('Inplane force [kN/m]')
    axs1[2,1].set_ylabel('Out-of-plane force [kN/m]')
    axs1[0,0].legend(['CS','HS2'])
    axs1_2[0][0].plot(bem_data[1:-1,0],1e3*(np.interp(bem_data[1:-1,0],defl_data[:,0],defl_data[:,4])-hs2_ind[:,9]),'.-g')
    axs1_2[1][0].plot(bem_data[1:-1,0],1e3*(np.interp(bem_data[1:-1,0],defl_data[:,0],defl_data[:,5])-hs2_ind[:,10]),'.-g')
    axs1_2[2][0].plot(bem_data[1:-1,0],1e3*(np.interp(bem_data[1:-1,0],defl_data[:,0],defl_data[:,6])-hs2_ind[:,11]),'.-g')
    axs1_2[0][1].plot(bem_data[1:-1,0],np.interp(bem_data[1:-1,0],defl_data[:,0],defl_data[:,9])-np.degrees(hs2_ind[:,34]*hs2_ind[:,37]),'.-g')
    axs1_2[1][1].plot(bem_data[1:-1,0],bem_data[1:-1,5]-hs2_ind[:,6],'.-g')
    axs1_2[2][1].plot(bem_data[1:-1,0],bem_data[1:-1,6]-hs2_ind[:,7],'.-g')
    axs1_2[0][0].set_ylabel('Diff CS - HS2 [mm]')
    axs1_2[1][0].set_ylabel('Diff CS - HS2 [mm]')
    axs1_2[2][0].set_ylabel('Diff CS - HS2 [mm]')
    axs1_2[0][1].set_ylabel('Diff CS - HS2 [deg]')
    axs1_2[1][1].set_ylabel('Diff CS - HS2 [N/m]')
    axs1_2[2][1].set_ylabel('Diff CS - HS2 [N/m]')
    axs1[2,0].set_xlabel('Spanwise coordinate [m]')
    axs1[2,1].set_xlabel('Spanwise coordinate [m]')
    plt.tight_layout()
    # plt.show()
    fig1.savefig(os.path.join(work_folder, res_folder + '/deflforce{:d}ms.png'.format(int(wsp))),dpi=300)

    plt.close('all')

    veltri = models[iops].rotors[0].blades[0].states_and_forces()

    # Plot velocity triangles
    fig2=plt.figure(figsize=[12,6])
    axs2=fig2.subplots(3,2)
    axs2_2 = []
    for i in range(3):
        axdummy = []
        for j in range(2):
            axdummy.append(axs2[i,j].twinx())
        axs2_2.append(axdummy)
    axs2[0,0].plot(bem_data[:,0],bem_data[:,8],'.-', hs2_ind[:,0],hs2_ind[:, 16],'.-')
    axs2[1,0].plot(bem_data[:,0],bem_data[:,9],'.-', hs2_ind[:,0],hs2_ind[:, 17],'.-')
    axs2[2,0].plot(bem_data[:,0],bem_data[:,3],'.-' ,hs2_ind[:,0],np.degrees(hs2_ind[:, 4]),'.-')
    axs2[0,1].plot(bem_data[:,0],bem_data[:,4],'.-' ,hs2_ind[:,0],hs2_ind[:, 5],'.-')
    axs2[1,1].plot(bem_data[:,0],bem_data[:,1],'.-' ,hs2_ind[:,0],hs2_ind[:, 1],'.-')
    axs2[2,1].plot(bem_data[:,0],bem_data[:,2],'.-' ,hs2_ind[:,0],hs2_ind[:, 2],'.-')
    axs2[0,0].set_ylabel('CL [-]')
    axs2[1,0].set_ylabel('CD [-]')
    axs2[2,0].set_ylabel('AoA [deg]')
    axs2[0,1].set_ylabel('Relative speed [m/s]')
    axs2[1,1].set_ylabel('a [-]')
    axs2[2,1].set_ylabel('ap [-]')
    axs2[0,0].legend(['CS','HS2'])
    axs2_2[0][0].plot(bem_data[1:-1,0],bem_data[1:-1,8]-hs2_ind[:, 16],'.-g')
    axs2_2[1][0].plot(bem_data[1:-1,0],bem_data[1:-1,9]-hs2_ind[:, 17],'.-g')
    axs2_2[2][0].plot(bem_data[1:-1,0],bem_data[1:-1,3]-np.degrees(hs2_ind[:, 4]),'.-g')
    axs2_2[0][1].plot(bem_data[1:-1,0],bem_data[1:-1,4]-hs2_ind[:, 5],'.-g')
    axs2_2[1][1].plot(bem_data[1:-1,0],bem_data[1:-1,1]-hs2_ind[:, 1],'.-g')
    axs2_2[2][1].plot(bem_data[1:-1,0],bem_data[1:-1,2]-hs2_ind[:, 2],'.-g')
    axs2_2[0][0].set_ylabel('Diff CS - HS2 [-]')
    axs2_2[1][0].set_ylabel('Diff CS - HS2 [-]')
    axs2_2[2][0].set_ylabel('Diff CS - HS2 [deg]')
    axs2_2[0][1].set_ylabel('Diff CS - HS2 [m/s]')
    axs2_2[1][1].set_ylabel('Diff CS - HS2 [-]')
    axs2_2[2][1].set_ylabel('Diff CS - HS2 [-]')
    axs2[2,0].set_xlabel('Spanwise position [m]')
    axs2[2,1].set_xlabel('Spanwise position [m]')
    plt.tight_layout()
    # plt.show()
    fig2.savefig(os.path.join(work_folder,res_folder + '/veltri{:d}ms.png'.format(int(wsp))),dpi=300)

    plt.close('all')



hs2_pwr = np.loadtxt(os.path.join(work_folder,'HAWCStab2_results/hs2.pwr'),skiprows = 1)

# Plot power curve
fig = plt.figure()
ax = fig.gca()
ax2 = ax.twinx()

ax.plot(ops[:,0],pwr*1e-3,'o-',hs2_pwr[:,0],hs2_pwr[:,1],'o-')
ax2.plot(ops[:,0],pwr*1e-3-hs2_pwr[:,1],'o-g')

ax.set_ylabel('Power [kW]')
ax.legend(['CS','HS2'])
ax2.set_ylabel('Diff CS - HS2 [kW]')
ax.set_xlabel('Wind speed [m/s]')

ax.set_xlim([ops[0,0],ops[-1,0]])

plt.tight_layout()
#plt.show()
fig.savefig(os.path.join(work_folder,res_folder + '/power.png'),dpi=300)

# Plot thrust curve
fig = plt.figure()
ax = fig.gca()
ax2 = ax.twinx()

ax.plot(ops[:,0],thr*1e-3,'o-',hs2_pwr[:,0],hs2_pwr[:,2],'o-')
ax2.plot(ops[:,0],thr*1e-3-hs2_pwr[:,2],'o-g')

ax.set_ylabel('Thrust [kN]')
ax.legend(['CS','HS2'])
ax2.set_ylabel('Diff CS - HS2 [kN]')
ax.set_xlabel('Wind speed [m/s]')

ax.set_xlim([ops[0,0],ops[-1,0]])

plt.tight_layout()
#plt.show()
fig.savefig(os.path.join(work_folder,res_folder + '/thrust.png'),dpi=300)




pass

"""